import networkx as nx
import matplotlib.pyplot as plt
import pulp
import random
import numpy as np

# ---------------------------
# 1. 随机管网生成模块
# ---------------------------
class GasNetworkGenerator:
    def __init__(self, n_nodes=20, n_sources=3, n_leaks=2):
        self.n_nodes = n_nodes
        self.n_sources = n_sources
        self.n_leaks = n_leaks
        self.G = nx.Graph()
        self.params = {
            'diameters': [100, 200, 300],
            'pressures': [0.4, 0.6, 0.8],
            'users': [50, 100, 200]
        }

    def _generate_topology(self):
        # 生成混合型拓扑结构
        G = nx.Graph()
        
        # 生成树状主干网
        tree = nx.random_tree(self.n_nodes)
        nx.add_path(G, list(tree.nodes()))
        
        # 添加环状连接
        for i in range(0, self.n_nodes-2, 3):
            if i+2 < self.n_nodes:
                G.add_edge(i, i+2, length=200)
        
        # 添加多源节点
        sources = random.sample(list(G.nodes), self.n_sources)
        nx.set_node_attributes(G, {n: 'source' for n in sources}, 'type')
        
        return G

    def _assign_attributes(self, G):
        # 随机分配属性
        for node in G.nodes:
            G.nodes[node]['diameter'] = random.choice(self.params['diameters'])
            G.nodes[node]['pressure'] = random.choice(self.params['pressures'])
            G.nodes[node]['users'] = random.choice(self.params['users'])
            G.nodes[node]['valve'] = random.random() < 0.3  # 30%概率有阀站
            
        # 随机生成泄漏点
        leaks = random.sample([n for n in G.nodes if G.nodes[n]['type']!='source'], 
                            self.n_leaks)
        nx.set_node_attributes(G, {n: 'leak' for n in leaks}, 'status')
        
        # 随机分配流量
        for u, v in G.edges:
            G.edges[u,v]['flow'] = random.randint(500, 2000)
            G.edges[u,v]['length'] = random.randint(100, 500)
            
        return G

    def generate(self):
        G = self._generate_topology()
        return self._assign_attributes(G)

# ---------------------------
# 2. 最小割隔离模块
# ---------------------------
def find_min_cut(G, leaks):
    # 创建流量网络副本
    flow_net = nx.DiGraph()
    
    # 添加反向边
    for u, v, data in G.edges(data=True):
        flow_net.add_edge(u, v, capacity=data['flow'])
        flow_net.add_edge(v, u, capacity=data['flow'])
    
    # 合并泄漏点为超级源
    super_source = 'super_source'
    flow_net.add_node(super_source)
    for leak in leaks:
        flow_net.add_edge(super_source, leak, capacity=np.inf)
    
    # 找到最小割
    min_cut = nx.minimum_cut(flow_net, super_source, 'super_sink')
    return min_cut[1]  # 返回被隔离的节点集

# ---------------------------
# 3. 多目标优化模块
# ---------------------------
def multiobjective_optimization(G, leaks, min_cut_nodes):
    # 创建线性规划模型
    model = pulp.LpProblem("Gas_Valve_Optimization", pulp.LpMinimize)
    
    # 决策变量：是否关闭阀站
    valves = [n for n in G.nodes if G.nodes[n]['valve']]
    x = pulp.LpVariable.dicts("valve", valves, cat='Binary')
    
    # 目标函数（加权求和）
    # 权重参数（可根据实际情况调整）
    weights = {
        'users': 0.4,
        'leak': 0.3,
        'time': 0.3
    }
    
    # 用户影响项
    user_cost = pulp.lpSum([
        G.nodes[n]['users'] * x[n] 
        for n in valves if G.nodes[n]['valve']
    ])
    
    # 泄漏控制项
    leak_cost = pulp.lpSum([
        1e5 * x[n] 
        for n in valves if n in min_cut_nodes
    ])
    
    # 时间成本项
    time_cost = pulp.lpSum([
        G.nodes[n]['pressure'] * x[n] 
        for n in valves if G.nodes[n]['valve']
    ])
    
    # 目标函数
    model += weights['users']*user_cost + weights['leak']*leak_cost + weights['time']*time_cost
    
    # 约束条件
    # 1. 必须隔离所有泄漏点
    for leak in leaks:
        model += pulp.lpSum([x[n] for n in nx.single_source_shortest_path(G, leak)]) >= 1
    
    # 2. 最小割节点必须关闭
    for node in min_cut_nodes:
        if node in valves:
            model += x[node] == 1
    
    # 3. 至少保留一个源点
    sources = [n for n in G.nodes if G.nodes[n]['type']=='source']
    model += pulp.lpSum([1 - x[n] for n in sources if n in valves]) >= 1
    
    # 求解
    model.solve()
    
    return model, x

# ---------------------------
# 4. 可视化模块
# ---------------------------
def visualize(G, model, x):
    pos = nx.spring_layout(G, k=0.5)
    
    # 绘制基础网络
    nx.draw_networkx_nodes(G, pos, 
                        node_color='lightgray',
                        node_size=500)
    
    # 标注属性
    labels = nx.get_node_attributes(G, 'status')
    nx.draw_networkx_labels(G, pos, labels=labels, font_size=8)
    
    # 标注阀站
    valves = [n for n in G.nodes if G.nodes[n]['valve']]
    nx.draw_networkx_nodes(G, pos, 
                        nodelist=valves,
                        node_color='lightblue',
                        node_size=600)
    
    # 标注关闭阀站
    closed = [n for n in valves if pulp.value(x[n]) == 1]
    nx.draw_networkx_nodes(G, pos, 
                        nodelist=closed,
                        node_color='red',
                        node_size=700)
    
    # 标注泄漏点
    leaks = [n for n in G.nodes if G.nodes[n]['status']=='leak']
    nx.draw_networkx_nodes(G, pos, 
                        nodelist=leaks,
                        node_color='yellow',
                        node_size=800)
    
    plt.title("Optimal Valve Closure Plan")
    plt.show()

# ---------------------------
# 5. 主程序
# ---------------------------
if __name__ == "__main__":
    # 生成管网
    generator = GasNetworkGenerator(n_nodes=20, n_sources=3, n_leaks=2)
    G = generator.generate()
    
    # 找到泄漏点
    leaks = [n for n in G.nodes if G.nodes[n]['status']=='leak']
    
    # 最小割隔离
    min_cut_nodes = find_min_cut(G, leaks)[1]
    print(f"Minimum Cut Nodes: {min_cut_nodes}")
    
    # 多目标优化
    model, x = multiobjective_optimization(G, leaks, min_cut_nodes)
    
    # 输出结果
    print("\nOptimization Results:")
    print(f"Status: {pulp.LpStatus[model.status]}")
    print(f"Objective Value: {pulp.value(model.objective)}")
    
    # 可视化
    visualize(G, model, x)
